!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief Calculate the electrostatic energy by the Smooth Particle Ewald method
!> \par History
!>      JGH (03-May-2001) : first correctly working version
!> \author JGH (21-Mar-2001)
! **************************************************************************************************
MODULE spme

   USE atomic_kind_types,               ONLY: get_atomic_kind
   USE atprop_types,                    ONLY: atprop_type
   USE bibliography,                    ONLY: Essmann1995,&
                                              cite_reference
   USE cell_types,                      ONLY: cell_type
   USE dgs,                             ONLY: dg_sum_patch,&
                                              dg_sum_patch_force_1d,&
                                              dg_sum_patch_force_3d
   USE ewald_environment_types,         ONLY: ewald_env_get,&
                                              ewald_environment_type
   USE ewald_pw_types,                  ONLY: ewald_pw_get,&
                                              ewald_pw_type
   USE kinds,                           ONLY: dp
   USE mathconstants,                   ONLY: fourpi
   USE message_passing,                 ONLY: mp_comm_type
   USE particle_types,                  ONLY: particle_type
   USE pme_tools,                       ONLY: get_center,&
                                              set_list
   USE pw_grid_types,                   ONLY: pw_grid_type
   USE pw_grids,                        ONLY: get_pw_grid_info
   USE pw_methods,                      ONLY: pw_copy,&
                                              pw_derive,&
                                              pw_integral_a2b,&
                                              pw_multiply_with,&
                                              pw_transfer
   USE pw_poisson_methods,              ONLY: pw_poisson_rebuild,&
                                              pw_poisson_solve
   USE pw_poisson_types,                ONLY: greens_fn_type,&
                                              pw_poisson_type
   USE pw_pool_types,                   ONLY: pw_pool_type
   USE pw_types,                        ONLY: pw_c1d_gs_type,&
                                              pw_r3d_rs_type
   USE realspace_grid_types,            ONLY: realspace_grid_desc_type,&
                                              realspace_grid_type,&
                                              rs_grid_create,&
                                              rs_grid_release,&
                                              rs_grid_set_box,&
                                              rs_grid_zero,&
                                              transfer_pw2rs,&
                                              transfer_rs2pw
   USE shell_potential_types,           ONLY: shell_kind_type
#include "./base/base_uses.f90"

   IMPLICIT NONE

   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'spme'

   PRIVATE
   PUBLIC :: spme_evaluate, spme_potential, spme_forces, get_patch

   INTERFACE get_patch
      MODULE PROCEDURE get_patch_a, get_patch_b
   END INTERFACE

CONTAINS

! **************************************************************************************************
!> \brief ...
!> \param ewald_env ...
!> \param ewald_pw ...
!> \param box ...
!> \param particle_set ...
!> \param fg_coulomb ...
!> \param vg_coulomb ...
!> \param pv_g ...
!> \param shell_particle_set ...
!> \param core_particle_set ...
!> \param fgshell_coulomb ...
!> \param fgcore_coulomb ...
!> \param use_virial ...
!> \param charges ...
!> \param atprop ...
!> \par History
!>      JGH (03-May-2001) : SPME with charge definition
!> \author JGH (21-Mar-2001)
! **************************************************************************************************
   SUBROUTINE spme_evaluate(ewald_env, ewald_pw, box, particle_set, &
                            fg_coulomb, vg_coulomb, pv_g, shell_particle_set, core_particle_set, &
                            fgshell_coulomb, fgcore_coulomb, use_virial, charges, atprop)

      TYPE(ewald_environment_type), POINTER              :: ewald_env
      TYPE(ewald_pw_type), POINTER                       :: ewald_pw
      TYPE(cell_type), POINTER                           :: box
      TYPE(particle_type), DIMENSION(:), INTENT(IN)      :: particle_set
      REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT)        :: fg_coulomb
      REAL(KIND=dp), INTENT(OUT)                         :: vg_coulomb
      REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT)        :: pv_g
      TYPE(particle_type), DIMENSION(:), OPTIONAL, &
         POINTER                                         :: shell_particle_set, core_particle_set
      REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT), &
         OPTIONAL                                        :: fgshell_coulomb, fgcore_coulomb
      LOGICAL, INTENT(IN)                                :: use_virial
      REAL(KIND=dp), DIMENSION(:), OPTIONAL, POINTER     :: charges
      TYPE(atprop_type), POINTER                         :: atprop

      CHARACTER(len=*), PARAMETER                        :: routineN = 'spme_evaluate'

      INTEGER                                            :: handle, i, ig, ipart, j, n, ncore, &
                                                            npart, nshell, o_spline, p1, p1_shell
      INTEGER, ALLOCATABLE, DIMENSION(:, :)              :: center, core_center, shell_center
      INTEGER, DIMENSION(3)                              :: nd, npts
      LOGICAL                                            :: do_shell
      REAL(KIND=dp)                                      :: alpha, dvols, fat1, ffa, ffb
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: core_delta, delta, shell_delta
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: rhos
      REAL(KIND=dp), DIMENSION(3)                        :: fat
      REAL(KIND=dp), DIMENSION(3, 3)                     :: f_stress, h_stress
      TYPE(greens_fn_type), POINTER                      :: green
      TYPE(mp_comm_type)                                 :: group
      TYPE(pw_c1d_gs_type), DIMENSION(3)                 :: dphi_g
      TYPE(pw_c1d_gs_type), POINTER                      :: phi_g, rhob_g
      TYPE(pw_grid_type), POINTER                        :: grid_spme
      TYPE(pw_poisson_type), POINTER                     :: poisson_env
      TYPE(pw_pool_type), POINTER                        :: pw_pool
      TYPE(pw_r3d_rs_type), POINTER                      :: rhob_r
      TYPE(realspace_grid_desc_type), POINTER            :: rs_desc
      TYPE(realspace_grid_type), DIMENSION(3)            :: drpot
      TYPE(realspace_grid_type), POINTER                 :: rden, rpot

      NULLIFY (grid_spme, green, poisson_env, phi_g, rhob_g, rhob_r, pw_pool, rden, rpot)
      CALL timeset(routineN, handle)
      CALL cite_reference(Essmann1995)

      !-------------- INITIALISATION ---------------------
      CALL ewald_env_get(ewald_env, alpha=alpha, o_spline=o_spline, &
                         group=group)
      CALL ewald_pw_get(ewald_pw, pw_big_pool=pw_pool, rs_desc=rs_desc, &
                        poisson_env=poisson_env)
      CALL pw_poisson_rebuild(poisson_env)
      green => poisson_env%green_fft
      grid_spme => pw_pool%pw_grid

      npart = SIZE(particle_set)

      CALL get_pw_grid_info(grid_spme, npts=npts, dvol=dvols)

      n = o_spline
      ALLOCATE (rhos(n, n, n))
      ALLOCATE (rden)
      CALL rs_grid_create(rden, rs_desc)
      CALL rs_grid_set_box(grid_spme, rs=rden)
      CALL rs_grid_zero(rden)

      ALLOCATE (center(3, npart), delta(3, npart))
      CALL get_center(particle_set, box, center, delta, npts, n)

      IF (PRESENT(shell_particle_set) .AND. (PRESENT(core_particle_set))) THEN
         CPASSERT(ASSOCIATED(shell_particle_set))
         CPASSERT(ASSOCIATED(core_particle_set))
         nshell = SIZE(shell_particle_set)
         ncore = SIZE(core_particle_set)
         CPASSERT(nshell == ncore)
         ALLOCATE (shell_center(3, nshell), shell_delta(3, nshell))
         CALL get_center(shell_particle_set, box, shell_center, shell_delta, npts, n)
         ALLOCATE (core_center(3, nshell), core_delta(3, nshell))
         CALL get_center(core_particle_set, box, core_center, core_delta, npts, n)
      END IF

      !-------------- DENSITY CALCULATION ----------------
      ipart = 0
      ! Particles
      DO
         CALL set_list(particle_set, npart, center, p1, rden, ipart)
         IF (p1 == 0) EXIT

         do_shell = (particle_set(p1)%shell_index /= 0)
         IF (do_shell) CYCLE
         ! calculate function on small boxes
         CALL get_patch(particle_set, delta, green, p1, rhos, is_core=.FALSE., &
                        is_shell=.FALSE., unit_charge=.FALSE., charges=charges)

         ! add boxes to real space grid (big box)
         CALL dg_sum_patch(rden, rhos, center(:, p1))
      END DO
      ! Shell-Model
      IF (PRESENT(shell_particle_set) .AND. PRESENT(core_particle_set)) THEN
         ipart = 0
         DO
            ! calculate function on small boxes
            CALL set_list(shell_particle_set, nshell, shell_center, p1_shell, &
                          rden, ipart)
            IF (p1_shell == 0) EXIT
            CALL get_patch(shell_particle_set, shell_delta, green, p1_shell, rhos, &
                           is_core=.FALSE., is_shell=.TRUE., unit_charge=.FALSE.)

            ! add boxes to real space grid (big box)
            CALL dg_sum_patch(rden, rhos, shell_center(:, p1_shell))
         END DO
         ipart = 0
         DO
            ! calculate function on small boxes
            CALL set_list(core_particle_set, nshell, core_center, p1_shell, &
                          rden, ipart)
            IF (p1_shell == 0) EXIT
            CALL get_patch(core_particle_set, core_delta, green, p1_shell, rhos, &
                           is_core=.TRUE., is_shell=.FALSE., unit_charge=.FALSE.)

            ! add boxes to real space grid (big box)
            CALL dg_sum_patch(rden, rhos, core_center(:, p1_shell))
         END DO
      END IF
      !----------- END OF DENSITY CALCULATION -------------

      NULLIFY (rhob_r)
      ALLOCATE (rhob_r)
      CALL pw_pool%create_pw(rhob_r)
      CALL transfer_rs2pw(rden, rhob_r)
      ! transform density to G space and add charge function
      NULLIFY (rhob_g)
      ALLOCATE (rhob_g)
      CALL pw_pool%create_pw(rhob_g)
      CALL pw_transfer(rhob_r, rhob_g)
      ! update charge function
      CALL pw_multiply_with(rhob_g, green%p3m_charge)

      !-------------- ELECTROSTATIC CALCULATION -----------
      ! allocate intermediate arrays
      DO i = 1, 3
         CALL pw_pool%create_pw(dphi_g(i))
      END DO
      NULLIFY (phi_g)
      ALLOCATE (phi_g)
      CALL pw_pool%create_pw(phi_g)
      CALL pw_poisson_solve(poisson_env, rhob_g, vg_coulomb, phi_g, dphi_g, &
                            h_stress)
      !---------- END OF ELECTROSTATIC CALCULATION --------

      ! Atomic Energy and Stress
      IF (atprop%energy .OR. atprop%stress) THEN
         ALLOCATE (rpot)
         CALL rs_grid_create(rpot, rs_desc)
         CALL rs_grid_set_box(grid_spme, rs=rpot)
         CALL pw_multiply_with(phi_g, green%p3m_charge)
         CALL pw_transfer(phi_g, rhob_r)
         CALL transfer_pw2rs(rpot, rhob_r)
         ipart = 0
         DO
            CALL set_list(particle_set, npart, center, p1, rden, ipart)
            IF (p1 == 0) EXIT
            do_shell = (particle_set(p1)%shell_index /= 0)
            IF (do_shell) CYCLE
            ! calculate function on small boxes
            CALL get_patch(particle_set, delta, green, p1, rhos, is_core=.FALSE., &
                           is_shell=.FALSE., unit_charge=.FALSE., charges=charges)
            ! integrate box and potential
            CALL dg_sum_patch_force_1d(rpot, rhos, center(:, p1), fat1)
            IF (atprop%energy) THEN
               atprop%atener(p1) = atprop%atener(p1) + 0.5_dp*fat1*dvols
            END IF
            IF (atprop%stress) THEN
               atprop%atstress(1, 1, p1) = atprop%atstress(1, 1, p1) + 0.5_dp*fat1*dvols
               atprop%atstress(2, 2, p1) = atprop%atstress(2, 2, p1) + 0.5_dp*fat1*dvols
               atprop%atstress(3, 3, p1) = atprop%atstress(3, 3, p1) + 0.5_dp*fat1*dvols
            END IF
         END DO
         ! Core-shell model
         IF (PRESENT(shell_particle_set) .AND. PRESENT(core_particle_set)) THEN
            ipart = 0
            DO
               CALL set_list(shell_particle_set, nshell, shell_center, p1_shell, rden, ipart)
               IF (p1_shell == 0) EXIT
               CALL get_patch(shell_particle_set, shell_delta, green, p1_shell, rhos, &
                              is_core=.FALSE., is_shell=.TRUE., unit_charge=.FALSE.)
               CALL dg_sum_patch_force_1d(rpot, rhos, shell_center(:, p1_shell), fat1)
               p1 = shell_particle_set(p1_shell)%atom_index
               IF (atprop%energy) THEN
                  atprop%atener(p1) = atprop%atener(p1) + 0.5_dp*fat1*dvols
               END IF
            END DO
            ipart = 0
            DO
               CALL set_list(core_particle_set, nshell, core_center, p1_shell, rden, ipart)
               IF (p1_shell == 0) EXIT
               CALL get_patch(core_particle_set, core_delta, green, p1_shell, rhos, &
                              is_core=.TRUE., is_shell=.FALSE., unit_charge=.FALSE.)
               CALL dg_sum_patch_force_1d(rpot, rhos, core_center(:, p1_shell), fat1)
               p1 = core_particle_set(p1_shell)%atom_index
               IF (atprop%energy) THEN
                  atprop%atener(p1) = atprop%atener(p1) + 0.5_dp*fat1*dvols
               END IF
            END DO
         END IF
         IF (atprop%stress) THEN
            ffa = (0.5_dp/alpha)**2
            ffb = 1.0_dp/fourpi
            DO i = 1, 3
               DO ig = grid_spme%first_gne0, grid_spme%ngpts_cut_local
                  phi_g%array(ig) = ffb*dphi_g(i)%array(ig)*(ffa*grid_spme%gsq(ig) + 1.0_dp)
                  phi_g%array(ig) = phi_g%array(ig)*poisson_env%green_fft%influence_fn%array(ig)
               END DO
               IF (grid_spme%have_g0) phi_g%array(1) = 0.0_dp
               DO j = 1, i
                  nd = 0
                  nd(j) = 1
                  CALL pw_copy(phi_g, rhob_g)
                  CALL pw_derive(rhob_g, nd)
                  CALL pw_multiply_with(rhob_g, green%p3m_charge)
                  CALL pw_transfer(rhob_g, rhob_r)
                  CALL transfer_pw2rs(rpot, rhob_r)

                  ipart = 0
                  DO
                     CALL set_list(particle_set, npart, center, p1, rden, ipart)
                     IF (p1 == 0) EXIT
                     ! calculate function on small boxes
                     CALL get_patch(particle_set, delta, green, p1, rhos, &
                                    is_core=.FALSE., is_shell=.FALSE., unit_charge=.FALSE., charges=charges)
                     ! integrate box and potential
                     CALL dg_sum_patch_force_1d(rpot, rhos, center(:, p1), fat1)
                     atprop%atstress(i, j, p1) = atprop%atstress(i, j, p1) + fat1*dvols
                     IF (i /= j) atprop%atstress(j, i, p1) = atprop%atstress(j, i, p1) + fat1*dvols
                  END DO

               END DO
            END DO
         END IF
         CALL rs_grid_release(rpot)
         DEALLOCATE (rpot)
      END IF

      CALL pw_pool%give_back_pw(phi_g)
      CALL pw_pool%give_back_pw(rhob_g)
      DEALLOCATE (phi_g, rhob_g)

      !------------- STRESS TENSOR CALCULATION ------------
      IF (use_virial) THEN
         DO i = 1, 3
            DO j = i, 3
               f_stress(i, j) = pw_integral_a2b(dphi_g(i), dphi_g(j))
               f_stress(j, i) = f_stress(i, j)
            END DO
         END DO
         ffa = (1.0_dp/fourpi)*(0.5_dp/alpha)**2
         f_stress = -ffa*f_stress
         pv_g = h_stress + f_stress
      END IF
      !--------END OF STRESS TENSOR CALCULATION -----------
      ! move derivative of potential to real space grid and
      ! multiply by charge function in g-space
      DO i = 1, 3
         CALL rs_grid_create(drpot(i), rs_desc)
         CALL rs_grid_set_box(grid_spme, rs=drpot(i))
         CALL pw_multiply_with(dphi_g(i), green%p3m_charge)
         CALL pw_transfer(dphi_g(i), rhob_r)
         CALL pw_pool%give_back_pw(dphi_g(i))
         CALL transfer_pw2rs(drpot(i), rhob_r)
      END DO

      CALL pw_pool%give_back_pw(rhob_r)
      DEALLOCATE (rhob_r)
      !----------------- FORCE CALCULATION ----------------
      ! initialize the forces
      fg_coulomb = 0.0_dp
      ! Particles
      ipart = 0
      DO
         CALL set_list(particle_set, npart, center, p1, rden, ipart)
         IF (p1 == 0) EXIT

         do_shell = (particle_set(p1)%shell_index /= 0)
         IF (do_shell) CYCLE
         ! calculate function on small boxes
         CALL get_patch(particle_set, delta, green, p1, rhos, is_core=.FALSE., &
                        is_shell=.FALSE., unit_charge=.FALSE., charges=charges)

         ! add boxes to real space grid (big box)
         CALL dg_sum_patch_force_3d(drpot, rhos, center(:, p1), fat)
         fg_coulomb(1, p1) = fg_coulomb(1, p1) - fat(1)*dvols
         fg_coulomb(2, p1) = fg_coulomb(2, p1) - fat(2)*dvols
         fg_coulomb(3, p1) = fg_coulomb(3, p1) - fat(3)*dvols
      END DO
      ! Shell-Model
      IF (PRESENT(shell_particle_set) .AND. (PRESENT(core_particle_set))) THEN
         IF (PRESENT(fgshell_coulomb)) THEN
            ipart = 0
            fgshell_coulomb = 0.0_dp
            DO
               ! calculate function on small boxes
               CALL set_list(shell_particle_set, nshell, shell_center, p1_shell, &
                             rden, ipart)
               IF (p1_shell == 0) EXIT

               CALL get_patch(shell_particle_set, shell_delta, green, &
                              p1_shell, rhos, is_core=.FALSE., is_shell=.TRUE., unit_charge=.FALSE.)

               ! add boxes to real space grid (big box)
               CALL dg_sum_patch_force_3d(drpot, rhos, shell_center(:, p1_shell), fat)
               fgshell_coulomb(1, p1_shell) = fgshell_coulomb(1, p1_shell) - fat(1)*dvols
               fgshell_coulomb(2, p1_shell) = fgshell_coulomb(2, p1_shell) - fat(2)*dvols
               fgshell_coulomb(3, p1_shell) = fgshell_coulomb(3, p1_shell) - fat(3)*dvols

            END DO
         END IF
         IF (PRESENT(fgcore_coulomb)) THEN
            ipart = 0
            fgcore_coulomb = 0.0_dp
            DO
               ! calculate function on small boxes
               CALL set_list(core_particle_set, nshell, core_center, p1_shell, &
                             rden, ipart)
               IF (p1_shell == 0) EXIT

               CALL get_patch(core_particle_set, core_delta, green, &
                              p1_shell, rhos, is_core=.TRUE., is_shell=.FALSE., unit_charge=.FALSE.)

               ! add boxes to real space grid (big box)
               CALL dg_sum_patch_force_3d(drpot, rhos, core_center(:, p1_shell), fat)
               fgcore_coulomb(1, p1_shell) = fgcore_coulomb(1, p1_shell) - fat(1)*dvols
               fgcore_coulomb(2, p1_shell) = fgcore_coulomb(2, p1_shell) - fat(2)*dvols
               fgcore_coulomb(3, p1_shell) = fgcore_coulomb(3, p1_shell) - fat(3)*dvols
            END DO
         END IF

      END IF
      !--------------END OF FORCE CALCULATION -------------

      !------------------CLEANING UP ----------------------
      CALL rs_grid_release(rden)
      DEALLOCATE (rden)
      DO i = 1, 3
         CALL rs_grid_release(drpot(i))
      END DO

      DEALLOCATE (rhos)
      DEALLOCATE (center, delta)
      IF (ALLOCATED(shell_center)) THEN
         DEALLOCATE (shell_center, shell_delta)
      END IF
      IF (ALLOCATED(core_center)) THEN
         DEALLOCATE (core_center, core_delta)
      END IF
      CALL timestop(handle)

   END SUBROUTINE spme_evaluate

! **************************************************************************************************
!> \brief Calculate the electrostatic potential from particles A (charge A) at
!>        positions of particles B
!> \param ewald_env ...
!> \param ewald_pw ...
!> \param box ...
!> \param particle_set_a ...
!> \param charges_a ...
!> \param particle_set_b ...
!> \param potential ...
!> \par History
!>      JGH (23-Oct-2014) : adapted from SPME evaluate
!> \author JGH (23-Oct-2014)
! **************************************************************************************************
   SUBROUTINE spme_potential(ewald_env, ewald_pw, box, particle_set_a, charges_a, &
                             particle_set_b, potential)

      TYPE(ewald_environment_type), POINTER              :: ewald_env
      TYPE(ewald_pw_type), POINTER                       :: ewald_pw
      TYPE(cell_type), POINTER                           :: box
      TYPE(particle_type), DIMENSION(:), INTENT(IN)      :: particle_set_a
      REAL(KIND=dp), DIMENSION(:), POINTER               :: charges_a
      TYPE(particle_type), DIMENSION(:), INTENT(IN)      :: particle_set_b
      REAL(KIND=dp), DIMENSION(:), INTENT(INOUT)         :: potential

      CHARACTER(len=*), PARAMETER                        :: routineN = 'spme_potential'

      INTEGER                                            :: handle, ipart, n, npart_a, npart_b, &
                                                            o_spline, p1
      INTEGER, ALLOCATABLE, DIMENSION(:, :)              :: center
      INTEGER, DIMENSION(3)                              :: npts
      REAL(KIND=dp)                                      :: alpha, dvols, fat1
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: delta
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: rhos
      TYPE(greens_fn_type), POINTER                      :: green
      TYPE(mp_comm_type)                                 :: group
      TYPE(pw_c1d_gs_type), POINTER                      :: phi_g, rhob_g
      TYPE(pw_grid_type), POINTER                        :: grid_spme
      TYPE(pw_poisson_type), POINTER                     :: poisson_env
      TYPE(pw_pool_type), POINTER                        :: pw_pool
      TYPE(pw_r3d_rs_type), POINTER                      :: rhob_r
      TYPE(realspace_grid_desc_type), POINTER            :: rs_desc
      TYPE(realspace_grid_type), POINTER                 :: rden, rpot

      NULLIFY (grid_spme, green, poisson_env, phi_g, rhob_g, rhob_r, pw_pool, &
               rden, rpot)
      CALL timeset(routineN, handle)
      CALL cite_reference(Essmann1995)

      !-------------- INITIALISATION ---------------------
      CALL ewald_env_get(ewald_env, alpha=alpha, o_spline=o_spline, group=group)
      CALL ewald_pw_get(ewald_pw, pw_big_pool=pw_pool, rs_desc=rs_desc, poisson_env=poisson_env)
      CALL pw_poisson_rebuild(poisson_env)
      green => poisson_env%green_fft
      grid_spme => pw_pool%pw_grid

      npart_a = SIZE(particle_set_a)
      npart_b = SIZE(particle_set_b)

      CALL get_pw_grid_info(grid_spme, npts=npts, dvol=dvols)

      n = o_spline
      ALLOCATE (rhos(n, n, n))
      ALLOCATE (rden)
      CALL rs_grid_create(rden, rs_desc)
      CALL rs_grid_set_box(grid_spme, rs=rden)
      CALL rs_grid_zero(rden)

      ALLOCATE (center(3, npart_a), delta(3, npart_a))
      CALL get_center(particle_set_a, box, center, delta, npts, n)

      !-------------- DENSITY CALCULATION ----------------
      ipart = 0
      ! Particles
      DO
         CALL set_list(particle_set_a, npart_a, center, p1, rden, ipart)
         IF (p1 == 0) EXIT

         ! calculate function on small boxes
         CALL get_patch(particle_set_a, delta, green, p1, rhos, charges_a)

         ! add boxes to real space grid (big box)
         CALL dg_sum_patch(rden, rhos, center(:, p1))
      END DO
      DEALLOCATE (center, delta)
      !----------- END OF DENSITY CALCULATION -------------

      NULLIFY (rhob_r)
      ALLOCATE (rhob_r)
      CALL pw_pool%create_pw(rhob_r)
      CALL transfer_rs2pw(rden, rhob_r)
      ! transform density to G space and add charge function
      NULLIFY (rhob_g)
      ALLOCATE (rhob_g)
      CALL pw_pool%create_pw(rhob_g)
      CALL pw_transfer(rhob_r, rhob_g)
      ! update charge function
      CALL pw_multiply_with(rhob_g, green%p3m_charge)

      !-------------- ELECTROSTATIC CALCULATION -----------
      ! allocate intermediate arrays
      NULLIFY (phi_g)
      ALLOCATE (phi_g)
      CALL pw_pool%create_pw(phi_g)
      CALL pw_poisson_solve(poisson_env, density=rhob_g, vhartree=phi_g)
      !---------- END OF ELECTROSTATIC CALCULATION --------

      !-------------- POTENTAIL AT ATOMIC POSITIONS -------
      ALLOCATE (center(3, npart_b), delta(3, npart_b))
      CALL get_center(particle_set_b, box, center, delta, npts, n)
      rpot => rden
      CALL pw_multiply_with(phi_g, green%p3m_charge)
      CALL pw_transfer(phi_g, rhob_r)
      CALL transfer_pw2rs(rpot, rhob_r)
      ipart = 0
      DO
         CALL set_list(particle_set_b, npart_b, center, p1, rden, ipart)
         IF (p1 == 0) EXIT
         ! calculate function on small boxes
         CALL get_patch(particle_set_b, delta, green, p1, rhos, &
                        is_core=.FALSE., is_shell=.FALSE., unit_charge=.TRUE.)
         ! integrate box and potential
         CALL dg_sum_patch_force_1d(rpot, rhos, center(:, p1), fat1)
         potential(p1) = potential(p1) + fat1*dvols
      END DO

      !------------------CLEANING UP ----------------------
      CALL pw_pool%give_back_pw(phi_g)
      CALL pw_pool%give_back_pw(rhob_g)
      CALL pw_pool%give_back_pw(rhob_r)
      DEALLOCATE (phi_g, rhob_g, rhob_r)
      CALL rs_grid_release(rden)
      DEALLOCATE (rden)

      DEALLOCATE (rhos)
      DEALLOCATE (center, delta)
      CALL timestop(handle)

   END SUBROUTINE spme_potential

! **************************************************************************************************
!> \brief Calculate the forces on particles B for the electrostatic interaction
!>        betrween particles A and B
!> \param ewald_env ...
!> \param ewald_pw ...
!> \param box ...
!> \param particle_set_a ...
!> \param charges_a ...
!> \param particle_set_b ...
!> \param charges_b ...
!> \param forces_b ...
!> \par History
!>      JGH (27-Oct-2014) : adapted from SPME evaluate
!> \author JGH (23-Oct-2014)
! **************************************************************************************************
   SUBROUTINE spme_forces(ewald_env, ewald_pw, box, particle_set_a, charges_a, &
                          particle_set_b, charges_b, forces_b)

      TYPE(ewald_environment_type), POINTER              :: ewald_env
      TYPE(ewald_pw_type), POINTER                       :: ewald_pw
      TYPE(cell_type), POINTER                           :: box
      TYPE(particle_type), DIMENSION(:), INTENT(IN)      :: particle_set_a
      REAL(KIND=dp), DIMENSION(:), POINTER               :: charges_a
      TYPE(particle_type), DIMENSION(:), INTENT(IN)      :: particle_set_b
      REAL(KIND=dp), DIMENSION(:), POINTER               :: charges_b
      REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT)      :: forces_b

      CHARACTER(len=*), PARAMETER                        :: routineN = 'spme_forces'

      INTEGER                                            :: handle, i, ipart, n, npart_a, npart_b, &
                                                            o_spline, p1
      INTEGER, ALLOCATABLE, DIMENSION(:, :)              :: center
      INTEGER, DIMENSION(3)                              :: npts
      REAL(KIND=dp)                                      :: alpha, dvols
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: delta
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: rhos
      REAL(KIND=dp), DIMENSION(3)                        :: fat
      TYPE(greens_fn_type), POINTER                      :: green
      TYPE(mp_comm_type)                                 :: group
      TYPE(pw_c1d_gs_type), DIMENSION(3)                 :: dphi_g
      TYPE(pw_c1d_gs_type), POINTER                      :: phi_g, rhob_g
      TYPE(pw_grid_type), POINTER                        :: grid_spme
      TYPE(pw_poisson_type), POINTER                     :: poisson_env
      TYPE(pw_pool_type), POINTER                        :: pw_pool
      TYPE(pw_r3d_rs_type), POINTER                      :: rhob_r
      TYPE(realspace_grid_desc_type), POINTER            :: rs_desc
      TYPE(realspace_grid_type), DIMENSION(3)            :: drpot
      TYPE(realspace_grid_type), POINTER                 :: rden

      NULLIFY (grid_spme, green, poisson_env, phi_g, rhob_g, rhob_r, &
               pw_pool, rden)
      CALL timeset(routineN, handle)
      CALL cite_reference(Essmann1995)

      !-------------- INITIALISATION ---------------------
      CALL ewald_env_get(ewald_env, alpha=alpha, o_spline=o_spline, group=group)
      CALL ewald_pw_get(ewald_pw, pw_big_pool=pw_pool, rs_desc=rs_desc, poisson_env=poisson_env)
      CALL pw_poisson_rebuild(poisson_env)
      green => poisson_env%green_fft
      grid_spme => pw_pool%pw_grid

      npart_a = SIZE(particle_set_a)
      npart_b = SIZE(particle_set_b)

      CALL get_pw_grid_info(grid_spme, npts=npts, dvol=dvols)

      n = o_spline
      ALLOCATE (rhos(n, n, n))
      ALLOCATE (rden)
      CALL rs_grid_create(rden, rs_desc)
      CALL rs_grid_set_box(grid_spme, rs=rden)
      CALL rs_grid_zero(rden)

      ALLOCATE (center(3, npart_a), delta(3, npart_a))
      CALL get_center(particle_set_a, box, center, delta, npts, n)

      !-------------- DENSITY CALCULATION ----------------
      ipart = 0
      ! Particles
      DO
         CALL set_list(particle_set_a, npart_a, center, p1, rden, ipart)
         IF (p1 == 0) EXIT

         ! calculate function on small boxes
         CALL get_patch(particle_set_a, delta, green, p1, rhos, charges_a)

         ! add boxes to real space grid (big box)
         CALL dg_sum_patch(rden, rhos, center(:, p1))
      END DO
      DEALLOCATE (center, delta)
      !----------- END OF DENSITY CALCULATION -------------

      NULLIFY (rhob_r)
      ALLOCATE (rhob_r)
      CALL pw_pool%create_pw(rhob_r)
      CALL transfer_rs2pw(rden, rhob_r)
      ! transform density to G space and add charge function
      NULLIFY (rhob_g)
      ALLOCATE (rhob_g)
      CALL pw_pool%create_pw(rhob_g)
      CALL pw_transfer(rhob_r, rhob_g)
      ! update charge function
      CALL pw_multiply_with(rhob_g, green%p3m_charge)

      !-------------- ELECTROSTATIC CALCULATION -----------
      ! allocate intermediate arrays
      DO i = 1, 3
         CALL pw_pool%create_pw(dphi_g(i))
      END DO
      NULLIFY (phi_g)
      ALLOCATE (phi_g)
      CALL pw_pool%create_pw(phi_g)
      CALL pw_poisson_solve(poisson_env, density=rhob_g, vhartree=phi_g, &
                            dvhartree=dphi_g)
      !---------- END OF ELECTROSTATIC CALCULATION --------
      ! move derivative of potential to real space grid and
      ! multiply by charge function in g-space
      DO i = 1, 3
         CALL rs_grid_create(drpot(i), rs_desc)
         CALL rs_grid_set_box(grid_spme, rs=drpot(i))
         CALL pw_multiply_with(dphi_g(i), green%p3m_charge)
         CALL pw_transfer(dphi_g(i), rhob_r)
         CALL pw_pool%give_back_pw(dphi_g(i))
         CALL transfer_pw2rs(drpot(i), rhob_r)
      END DO
      !----------------- FORCE CALCULATION ----------------
      ALLOCATE (center(3, npart_b), delta(3, npart_b))
      CALL get_center(particle_set_b, box, center, delta, npts, n)
      ipart = 0
      DO
         CALL set_list(particle_set_b, npart_b, center, p1, rden, ipart)
         IF (p1 == 0) EXIT
         ! calculate function on small boxes
         CALL get_patch(particle_set_b, delta, green, p1, rhos, &
                        is_core=.FALSE., is_shell=.FALSE., unit_charge=.FALSE., charges=charges_b)
         ! add boxes to real space grid (big box)
         CALL dg_sum_patch_force_3d(drpot, rhos, center(:, p1), fat)
         forces_b(1, p1) = forces_b(1, p1) - fat(1)*dvols
         forces_b(2, p1) = forces_b(2, p1) - fat(2)*dvols
         forces_b(3, p1) = forces_b(3, p1) - fat(3)*dvols
      END DO
      !------------------CLEANING UP ----------------------
      DO i = 1, 3
         CALL rs_grid_release(drpot(i))
      END DO
      CALL pw_pool%give_back_pw(phi_g)
      CALL pw_pool%give_back_pw(rhob_g)
      CALL pw_pool%give_back_pw(rhob_r)
      DEALLOCATE (phi_g, rhob_g, rhob_r)
      CALL rs_grid_release(rden)
      DEALLOCATE (rden)

      DEALLOCATE (rhos)
      DEALLOCATE (center, delta)
      CALL timestop(handle)

   END SUBROUTINE spme_forces

! **************************************************************************************************
!> \brief Calculates local density in a small box
!> \param part ...
!> \param delta ...
!> \param green ...
!> \param p ...
!> \param rhos ...
!> \param is_core ...
!> \param is_shell ...
!> \param unit_charge ...
!> \param charges ...
!> \par History
!>      none
!> \author JGH (21-Mar-2001)
! **************************************************************************************************
   SUBROUTINE get_patch_a(part, delta, green, p, rhos, is_core, is_shell, &
                          unit_charge, charges)

      TYPE(particle_type), DIMENSION(:), INTENT(IN)      :: part
      REAL(KIND=dp), DIMENSION(:, :), INTENT(IN)         :: delta
      TYPE(greens_fn_type), INTENT(IN)                   :: green
      INTEGER, INTENT(IN)                                :: p
      REAL(KIND=dp), DIMENSION(:, :, :), INTENT(OUT)     :: rhos
      LOGICAL, INTENT(IN)                                :: is_core, is_shell, unit_charge
      REAL(KIND=dp), DIMENSION(:), OPTIONAL, POINTER     :: charges

      INTEGER                                            :: nbox
      LOGICAL                                            :: use_charge_array
      REAL(KIND=dp)                                      :: q
      REAL(KIND=dp), DIMENSION(3)                        :: r
      TYPE(shell_kind_type), POINTER                     :: shell

      NULLIFY (shell)
      use_charge_array = PRESENT(charges)
      IF (use_charge_array) use_charge_array = ASSOCIATED(charges)
      IF (is_core .AND. is_shell) THEN
         CPABORT("Shell-model: cannot be core and shell simultaneously")
      END IF

      nbox = SIZE(rhos, 1)
      r = part(p)%r
      q = 1.0_dp
      IF (.NOT. unit_charge) THEN
         IF (is_core) THEN
            CALL get_atomic_kind(atomic_kind=part(p)%atomic_kind, shell=shell)
            q = shell%charge_core
         ELSE IF (is_shell) THEN
            CALL get_atomic_kind(atomic_kind=part(p)%atomic_kind, shell=shell)
            q = shell%charge_shell
         ELSE
            CALL get_atomic_kind(atomic_kind=part(p)%atomic_kind, qeff=q)
         END IF
         IF (use_charge_array) q = charges(p)
      END IF
      CALL spme_get_patch(rhos, nbox, delta(:, p), q, green%p3m_coeff)

   END SUBROUTINE get_patch_a

! **************************************************************************************************
!> \brief Calculates local density in a small box
!> \param part ...
!> \param delta ...
!> \param green ...
!> \param p ...
!> \param rhos ...
!> \param charges ...
!> \par History
!>      none
!> \author JGH (21-Mar-2001)
! **************************************************************************************************
   SUBROUTINE get_patch_b(part, delta, green, p, rhos, charges)

      TYPE(particle_type), DIMENSION(:), INTENT(IN)      :: part
      REAL(KIND=dp), DIMENSION(:, :), INTENT(IN)         :: delta
      TYPE(greens_fn_type), INTENT(IN)                   :: green
      INTEGER, INTENT(IN)                                :: p
      REAL(KIND=dp), DIMENSION(:, :, :), INTENT(OUT)     :: rhos
      REAL(KIND=dp), DIMENSION(:), POINTER               :: charges

      INTEGER                                            :: nbox
      REAL(KIND=dp)                                      :: q
      REAL(KIND=dp), DIMENSION(3)                        :: r

      nbox = SIZE(rhos, 1)
      r = part(p)%r
      q = charges(p)
      CALL spme_get_patch(rhos, nbox, delta(:, p), q, green%p3m_coeff)

   END SUBROUTINE get_patch_b

! **************************************************************************************************
!> \brief Calculates SPME charge assignment
!> \param rhos ...
!> \param n ...
!> \param delta ...
!> \param q ...
!> \param coeff ...
!> \par History
!>      DG (29-Mar-2001) : code implemented
!> \author JGH (22-Mar-2001)
! **************************************************************************************************
   SUBROUTINE spme_get_patch(rhos, n, delta, q, coeff)

      REAL(KIND=dp), DIMENSION(:, :, :), INTENT(OUT)     :: rhos
      INTEGER, INTENT(IN)                                :: n
      REAL(KIND=dp), DIMENSION(3), INTENT(IN)            :: delta
      REAL(KIND=dp), INTENT(IN)                          :: q
      REAL(KIND=dp), DIMENSION(-(n-1):n-1, 0:n-1), &
         INTENT(IN)                                      :: coeff

      INTEGER, PARAMETER                                 :: nmax = 12

      INTEGER                                            :: i, i1, i2, i3, j, l
      REAL(KIND=dp)                                      :: r2, r3
      REAL(KIND=dp), DIMENSION(3, -nmax:nmax)            :: w_assign
      REAL(KIND=dp), DIMENSION(3, 0:nmax-1)              :: deltal
      REAL(KIND=dp), DIMENSION(3, 1:nmax)                :: f_assign

      IF (n > nmax) THEN
         CPABORT("nmax value too small")
      END IF
      ! calculate the assignment function values and
      ! the charges on the grid (small box)

      deltal(1, 0) = 1.0_dp
      deltal(2, 0) = 1.0_dp
      deltal(3, 0) = 1.0_dp
      DO l = 1, n - 1
         deltal(1, l) = deltal(1, l - 1)*delta(1)
         deltal(2, l) = deltal(2, l - 1)*delta(2)
         deltal(3, l) = deltal(3, l - 1)*delta(3)
      END DO

      w_assign = 0.0_dp
      DO j = -(n - 1), n - 1, 2
         DO l = 0, n - 1
            w_assign(1, j) = w_assign(1, j) + coeff(j, l)*deltal(1, l)
            w_assign(2, j) = w_assign(2, j) + coeff(j, l)*deltal(2, l)
            w_assign(3, j) = w_assign(3, j) + coeff(j, l)*deltal(3, l)
         END DO
      END DO
      DO i = 1, n
         j = n + 1 - 2*i
         f_assign(1, i) = w_assign(1, j)
         f_assign(2, i) = w_assign(2, j)
         f_assign(3, i) = w_assign(3, j)
      END DO

      DO i3 = 1, n
         r3 = q*f_assign(3, i3)
         DO i2 = 1, n
            r2 = r3*f_assign(2, i2)
            DO i1 = 1, n
               rhos(i1, i2, i3) = r2*f_assign(1, i1)
            END DO
         END DO
      END DO

   END SUBROUTINE spme_get_patch

END MODULE spme

